/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.common.io;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.Socket;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.Set;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.sss.common.EncodedBucket;
import org.sss.mapreduce.GroupID;
import org.sss.mapreduce.SssException;

/**
 * Client for original protocol storage server.
 */
public class StreamDB {
  private static Logger logger = LoggerFactory.getLogger(StreamDB.class);

  private static int MAX_RETRY = 5;

  private static int PROTOCOL_MAGIC = 0x55;
  private static int PROTOCOL_ENDMARK = -1;

  private static int PROTOCOL_COMMAND_GET    = 0x81;
  private static int PROTOCOL_COMMAND_PUT    = 0x82;
  private static int PROTOCOL_COMMAND_REMOVE = 0x83;
  private static int PROTOCOL_COMMAND_LIST   = 0x84;
  private static int PROTOCOL_COMMAND_DREAD  = 0x85;

  private Socket           socket;
  private DataOutputStream outputStream;
  private DataInputStream  inputStream;
  private boolean working = false;

  // for reconnecting
  private String           host;
  private int              port;

  public StreamDB(String host, int port) throws SssException {
    // delay connect
    this.host = host;
    this.port = port;
  }

  public void close() throws SssException {
    try {
      if (socket != null) {
        socket.close();
      }
    }
    catch (IOException e) {
      throw new SssException(e);
    }
  }

  public Iterable<byte[]> read(GroupID gid, int subhash, boolean destruct) throws SssException {
    int retry = 0;
    enter();
    Reader r = null;
    try {
      while (true) {
        try {
          reconnectIfInvalid();
          writeSimpleRequest(destruct ? PROTOCOL_COMMAND_DREAD : PROTOCOL_COMMAND_GET,
                             gid.toString(), subhash);
          readReplyCommon();

          r = new Reader();
          final Reader rd = r;
          return new Iterable<byte[]>() {
            @Override
            public Iterator<byte[]> iterator() {
              return rd;
            }
          };
        }
        catch (SssException e) {
          disable();
          ++retry;
          if (retry >= MAX_RETRY) {
            throw e;
          }
          else {
            logger.error("ignore exception: " + e);
            logger.error("try reconnect");
          }
        }
      }
    }
    finally {
      if (r == null) {
        disable();
        leave();
      }
    }
  }

  /**
   *
   * TODO: stream?
   * @return number of bytes of the written data.
   */
  public long write(GroupID gid, Iterable<EncodedBucket> data) throws SssException {
    int retry = 0;
    boolean suc = false;
    long nbytes = 0;
    enter();
    try {
      while (true) {
        try {
          try {
            reconnectIfInvalid();

            outputStream.write(PROTOCOL_MAGIC);
            outputStream.write(PROTOCOL_COMMAND_PUT);
            writeArg(gid.toString());

            int size = 0;
            for (EncodedBucket d: data) {
              size += 4 + d.key.length +
                4 + d.value.length;
              nbytes += d.key.length + d.value.length;
            }
            outputStream.writeInt(size);
            for (EncodedBucket d: data) {
              writeArg(d.key);
              writeArg(d.value);
            }
            writeEndMark();

            readSimpleReply(PROTOCOL_COMMAND_PUT, 0);
            suc = true;
            return nbytes;
          }
          catch (IOException e) {
            throw new SssException(e);
          }
        }
        catch (SssException e) {
          disable();
          ++retry;
          if (retry >= MAX_RETRY) {
            throw e;
          }
          else {
            logger.error("ignore exception: " + e);
            logger.error("try reconnect");
          }
        }
      }
    }
    finally {
      if (!suc) disable();
      leave();
    }
  }

  public void delete(GroupID gid) throws SssException {
    simpleRPC(PROTOCOL_COMMAND_REMOVE, 0, gid.toString());
  }

  public Set<GroupID> listTupleGroups() throws SssException {
    try {
      List<byte[]> result = simpleRPC(PROTOCOL_COMMAND_LIST, -1);

      Set<GroupID> gids = new HashSet<GroupID>();

      for (byte[] b: result) {
        String str = new String(b, "UTF-8");
        gids.add(GroupID.createFromString(str));
      }
      return gids;
    }
    catch (UnsupportedEncodingException e) {
      throw new SssException(e);
    }
  }

  private class Reader implements Iterator<byte[]> {
    private byte[] record;
    private boolean end = false;;

    @Override
    public boolean hasNext() {
      if (record != null) {
        return true;
      }
      if (end) {
        return false;
      }
      try {
        record = readArg();
        if (record == null) {
          end = true;
          leave();
          return false;
        }
        return true;
      }
      catch (SssException e) {
        disable();
        leave();
        throw new RuntimeException(e);
      }
    }
    @Override
    public byte[] next() {
      if (!hasNext()) {
        throw new NoSuchElementException();
      }
      byte[] tmp = record;
      record = null;
      return tmp;
    }
    @Override
    public void remove() {
      throw new UnsupportedOperationException();
    }
  }

  private List<byte[]> simpleRPC(int func, int nExpectResults, Object... params) throws SssException {
    List<byte[]> r = null;
    int retry = 0;
    enter();
    try {
      while (true) {
        try {
          reconnectIfInvalid();
          writeSimpleRequest(func, params);
          r = readSimpleReply(func, nExpectResults);
          return r;
        }
        catch (SssException e) {
          disable();
          ++retry;
          if (retry >= MAX_RETRY) {
            throw e;
          }
          else {
            logger.error("ignore exception: " + e);
            logger.error("try reconnect");
          }
        }
      }
    }
    finally {
      if (r == null) disable();
      leave();
    }
  }

  private void writeSimpleRequest(int command, Object... args) throws SssException {
    try {
      outputStream.write(PROTOCOL_MAGIC);
      outputStream.write(command);
      outputStream.flush();

      for (Object arg: args) {
        writeArg(arg);
      }
      writeEndMark();
    }
    catch (IOException e) {
      throw new SssException(e);
    }
  }

  private List<byte[]> readSimpleReply(int func, int nExpectResults) throws SssException {
    readReplyCommon();
    List<byte[]> o = new ArrayList<byte[]>();
    byte[] v;

    while ((v = readArg()) != null) {
      o.add(v);
    }

    if (nExpectResults >= 0 && nExpectResults != o.size()) {
      throw new SssException(
          String.format("funccode %s: Expects that result has %d parameters, but %d parameters",
            func, nExpectResults, o.size()));
    }
    return o;
  }

  private void readReplyCommon() throws SssException {
    try {
      int magic = inputStream.read();
      if (magic != PROTOCOL_MAGIC) {
        throw new SssException("The head of protocol header is not magic number.");
      }
      int result = inputStream.read();
      if (result == 1) {
        return; // Success
      }
      else if (result == 0) {
        byte[] b = readArg();
        String msg = new String(b, "UTF-8");
        if (readArg() != null) {
          throw new SssException("ProtocolError");
        }
        throw new SssException("Operation failed: " + msg);
      }
      else if (result < 0) {
        throw new SssException("Unexpect EOF");
      }
      else {
        // TODO: disable this.
        throw new SssException("Unkown result: " + result);
      }
    }
    catch (IOException e) {
      // TODO: disable this.
      throw new SssException(e);
    }
  }

  private void writeArg(Object arg) throws SssException
  {
    try {
      if (arg instanceof String) {
        byte[] b = ((String)arg).getBytes("UTF-8");
        outputStream.writeInt(b.length);
        outputStream.write(b);
      }
      else if (arg instanceof byte[]) {
        byte[] b = (byte[])arg;
        outputStream.writeInt(b.length);
        outputStream.write(b);
      }
      else if (arg instanceof Integer) {
        outputStream.writeInt(4);
        outputStream.writeInt((Integer)arg);
      }
      else {
        throw new IllegalArgumentException("" + arg.getClass().getName());
      }
    }
    catch (IOException e) {
      throw new SssException(e);
    }
  }

  private void writeEndMark() throws SssException {
    try {
      outputStream.writeInt(PROTOCOL_ENDMARK);
      outputStream.flush();
    }
    catch (IOException e) {
      throw new SssException(e);
    }
  }

  private byte[] readArg() throws SssException {
    try {
      int size = inputStream.readInt();
      if (size == PROTOCOL_ENDMARK) {
        return null;
      }
      byte[] b = new byte[size];
      inputStream.readFully(b);
      return b;
    }
    catch (IOException e) {
      throw new SssException(e);
    }
  }


  private synchronized void enter() {
    while (working) {
      try {
        wait();
      }
      catch (InterruptedException e) {
      }
    }
    working = true;
  }

  private synchronized void leave() {
    working = false;
    notifyAll();
  }

  private synchronized void disable() {
    try {
      if (socket != null) {
        socket.close();
      }
    }
    catch (IOException e) {
      logger.error("close failed:" + e);
    }
    socket = null;
    outputStream = null;
    inputStream = null;
  }

  private synchronized void reconnectIfInvalid() throws SssException {
    boolean suc = false;
    if (socket == null) {
      try {
        this.socket       = new Socket(host, port);
        // TODO: buffering?
        this.outputStream = new DataOutputStream(new BufferedOutputStream(socket.getOutputStream(), 32 * 1024));
        this.inputStream  = new DataInputStream(new BufferedInputStream(socket.getInputStream(), 32 * 1024));
        // this.inputStream  = new DataInputStream(socket.getInputStream());
        suc = true;
      }
      catch (IOException e) {
        throw new SssException(e);
      }
      finally {
        if (!suc) {
          disable();
        }
      }
    }
  }
}
